odeint can be told that the Jacobian is banded. For a large system of differential equations that are known to be stiff, this can improve performance significantly.

As an example, we’ll solve the 1-D Gray-Scott partial differential equations using the method of lines [MOL]. The Gray-Scott equations for the functions \(u(x, t)\) and \(v(x, t)\) on the interval \(x \in [0, L]\) are

where \(D_u\) and \(D_v\) are the diffusion coefficients of the components \(u\) and \(v\), respectively, and \(f\) and \(k\) are constants. (For more information about the system, see http://groups.csail.mit.edu/mac/projects/amorphous/GrayScott/)

We’ll assume Neumann (i.e., “no flux”) boundary conditions:

To apply the method of lines, we discretize the \(x\) variable by defining the uniformly spaced grid of \(N\) points \(\left\{x_0, x_1, \ldots, x_{N-1}\right\}\), with \(x_0 = 0\) and \(x_{N-1} = L\). We define \(u_j(t) \equiv u(x_k, t)\) and \(v_j(t) \equiv v(x_k, t)\), and replace the \(x\) derivatives with finite differences. That is,

We then have a system of \(2N\) ordinary differential equations:

For convenience, the \((t)\) arguments have been dropped.

To enforce the boundary conditions, we introduce “ghost” points \(x_{-1}\) and \(x_N\), and define \(u_{-1}(t) \equiv u_1(t)\), \(u_N(t) \equiv u_{N-2}(t)\); \(v_{-1}(t)\) and \(v_N(t)\) are defined analogously.

Then

and

Our complete system of \(2N\) ordinary differential equations is (1) for \(k = 1, 2, \ldots, N-2\), along with (2) and (3).

We can now starting implementing this system in code. We must combine \(\{u_k\}\) and \(\{v_k\}\) into a single vector of length \(2N\). The two obvious choices are \(\{u_0, u_1, \ldots, u_{N-1}, v_0, v_1, \ldots, v_{N-1}\}\) and \(\{u_0, v_0, u_1, v_1, \ldots, u_{N-1}, v_{N-1}\}\). Mathematically, it does not matter, but the choice affects how efficiently odeint can solve the system. The reason is in how the order affects the pattern of the nonzero elements of the Jacobian matrix.

When the variables are ordered as \(\{u_0, u_1, \ldots, u_{N-1}, v_0, v_1, \ldots, v_{N-1}\}\), the pattern of nonzero elements of the Jacobian matrix is

The Jacobian pattern with variables interleaved as \(\{u_0, v_0, u_1, v_1, \ldots, u_{N-1}, v_{N-1}\}\) is

In both cases, there are just five nontrivial diagonals, but when the variables are interleaved, the bandwidth is much smaller. That is, the main diagonal and the two diagonals immediately above and the two immediately below the main diagonal are the nonzero diagonals. This is important, because the inputs mu and ml of odeint are the upper and lower bandwidths of the Jacobian matrix. When the variables are interleaved, mu and ml are 2. When the variables are stacked with \(\{v_k\}\) following \(\{u_k\}\), the upper and lower bandwidths are \(N\).

With that decision made, we can write the function that implements the system of differential equations.

First, we define the functions for the source and reaction terms of the system:

Next, we define the function that computes the right-hand side of the system of differential equations:

We won’t implement a function to compute the Jacobian, but we will tell odeint that the Jacobian matrix is banded. This allows the underlying solver (LSODA) to avoid computing values that it knows are zero. For a large system, this improves the performance significantly, as demonstrated in the following ipython session.

First, we define the required inputs:

Time the computation without taking advantage of the banded structure of the Jacobian matrix:

Now set ml=2 and mu=2, so odeint knows that the Jacobian matrix is banded:

That is quite a bit faster!

Let’s ensure that they have computed the same result: